#include <iostream>
#include<opencv2/opencv.hpp>

#include<math.h>
#include "yolov8_seg_onnx.h"
#include "crypto.h"

#include<cmdline.h>
#include<time.h>
//#define  VIDEO_OPENCV //if define, use opencv for video.

using namespace std;
using namespace cv;

void processRecognitionResults(const std::vector<OutputParams>& result, const cv::Mat& img, const std::string& save_path,
                               std::vector<OutputParams>& oldman_result, std::vector<OutputParams>& bed_result, 
                               std::vector<OutputParams>& nurse_result_max)
{
	// 检查输入参数的有效性
    if (result.empty()) {
        throw std::runtime_error("识别结果为空");
    }
    if (img.empty()) {
        throw std::runtime_error("输入图像为空");
    }
    if (save_path.empty()) {
        throw std::runtime_error("保存路径为空或无效");
    }


	std::vector<OutputParams> nurse_result;
    // 过滤ID和置信度
    for (const OutputParams& res : result) {
        if (res.id == 0 && res.confidence > 0.3) {
            oldman_result.push_back(res);
        } else if (res.id == 1 && res.confidence > 0.5) {
            bed_result.push_back(res);
        } else if (res.id == 2 && res.confidence > 0.5) {
            nurse_result.push_back(res);
        }
    }

	// 计算床的最大面积和位置
    if (bed_result.empty()) {
        throw std::runtime_error("未检测到床");
    }


    float max_area = 0;
	OutputParams bed_max_result;
    for (const auto& bed : bed_result) {
        float area = bed.box.width * bed.box.height;
        if (area > max_area) {
            max_area = area;
            bed_max_result = bed;
        }
    }

    // 计算所有护士与最大床之间的距离，过滤掉不在床范围内的护士
    for (const auto& nurse : nurse_result) {
        float bed_distance = sqrt(pow(bed_max_result.box.width, 2) + pow(bed_max_result.box.height, 2));
        float distance = sqrt(pow(nurse.box.x - bed_max_result.box.x, 2) + pow(nurse.box.y - bed_max_result.box.y, 2));
        if (distance < bed_distance / 2) {
            nurse_result_max.push_back(nurse);
        }
    }


    cv::Mat res = img.clone();
    cv::Mat res_mask = img.clone();

    // 绘制护工mask
    for (const auto& nurse : nurse_result_max) {
        res_mask(nurse.box).setTo(cv::Scalar(0, 220, 220), nurse.boxMask);
    }
    cv::addWeighted(res, 0.7, res_mask, 0.3, 0, res);

    // 绘制老人mask
    for (const auto& oldman : oldman_result) {
        res(oldman.box).setTo(cv::Scalar(60, 20, 220), oldman.boxMask);
    }

    // 绘制床mask
    for (const auto& bed : bed_result) {
        res(bed.box).setTo(cv::Scalar(51, 102, 0), bed.boxMask);
    }

     // 保存结果
    if (!cv::imwrite(save_path, res)) {
        throw std::runtime_error("保存结果失败");
    }
}

template<typename _Tp>
int yolov8_onnx(_Tp& task, cv::Mat& img, std::string& model_path, bool is_gpu, const std::string& save_path)
{

	if (task.ReadModel(model_path, is_gpu)) {
		std::cout << "read net ok!" << std::endl;
	}
	else {
		return -1;
	}
	//生成随机颜色
	std::vector<cv::Scalar> color = {cv::Scalar(60, 20, 220),cv::Scalar(51, 102, 0),cv::Scalar(0, 220, 220)};
	srand(time(0));

	/*********************老人、床、护士分类并过滤*********************/

	std::vector<OutputParams> result;
	std::vector<OutputParams> oldman_result;
	std::vector<OutputParams> bed_result;
	std::vector<OutputParams> nurse_result;

	if (task.OnnxDetect(img, result)) {

		/*********************老人、床、护士分类并过滤*********************/
		processRecognitionResults(result, img, save_path, oldman_result, bed_result, nurse_result);
	

		DrawPred(img, result, task._className, color);
	}
	else {
		std::cout << "Detect Failed!" << std::endl;
	}
	system("pause");
	return 0;
}

template<typename _Tp>
int yolov8_onnx_stream(_Tp& task, cv::Mat& img, std::vector<char>& modelBinary, bool is_gpu, const std::string &save_path)
{

	if (task.ReadModelFromStream(modelBinary, is_gpu)) {
		std::cout << "read net ok!" << std::endl;
	}
	else {
		return -1;
	}
	//生成随机颜色
	std::vector<cv::Scalar> color = {cv::Scalar(60, 20, 220),cv::Scalar(51, 102, 0),cv::Scalar(0, 220, 220)};
	srand(time(0));


	std::vector<OutputParams> result;
	
	std::vector<OutputParams> oldman_result;
	std::vector<OutputParams> bed_result;
	std::vector<OutputParams> nurse_result;

	if (task.OnnxDetect(img, result)) {

		/*********************老人、床、护士分类并过滤*********************/
		processRecognitionResults(result, img, save_path, oldman_result, bed_result, nurse_result);

		DrawPred(img, result, task._className, color);
	}
	else {
		std::cout << "Detect Failed!" << std::endl;
	}
	system("pause");
	return 0;
}


int main(int argc, char* argv[]) {

	cmdline::parser cmd;
    cmd.add<std::string>("model_path", 'm', "Path to YOLO onnx model.", false, "../models/oldman_detect.bin");
    cmd.add<std::string>("image", 'i', "Image source to be detected.", false, "../images/all.jpg");
	cmd.add<std::string>("output", 'o', "Path to output image.", false, "../result.jpg");
    cmd.add("gpu", '\0', "Inference on cuda device.");

    cmd.parse_check(argc, argv);

    const std::string img_path = cmd.get<std::string>("image");
    const std::string model_Path = cmd.get<std::string>("model_path");
	const std::string output_path = cmd.get<std::string>("output");
	const bool isGpu = cmd.exist("gpu");

    unsigned char key[32];
	unsigned char iv[16] = "bolitech";

	
// 模型加密
#ifdef MODEL_CRYPTO
    try {
		const std::string originFile = "/home/boli/project/code/oldman_detect/models/oldman_detect.onnx";
    	const std::string encryFile = "/home/boli/project/code/oldman_detect/models/oldman_detect.bin";
    	const std::string decryFile = "/home/boli/project/code/oldman_detect/models/best.onnx";
        // 根据用户选择调用加密或解密函数
        bool isEncrypt = true; // 假设用户选择加密
        if (isEncrypt) {
			#ifdef OLY_KEY
				generateKey(key);
			#else
				generateKeyAndIV(key, iv);
			#endif
            
            encryptFile(originFile, encryFile, key, iv);
        } else {
			#ifdef OLY_KEY
				readKey(key);
			#else
				readKeyAndIV(key, iv);
			#endif
            decryptFile(encryFile, decryFile, key, iv);
        }
    } catch (const exception& e) {
        cerr << "Error: " << e.what() << endl;
        return 1;
    }

#endif

// 加密模型推理
#ifdef MODEL_DECRYPTO

	// 图片和模型路径
	std::vector<char> modelBinary;

	try{
		#ifdef OLY_KEY
			readKey(key);
		#else
			readKeyAndIV(key, iv);
		#endif

		modelBinary = decryptFileToBinaryStream(model_Path, key, iv);

	}catch (const exception& e) {
        cerr << "Error: " << e.what() << endl;
        return 1;
	}


	// 加载图片
	cv::Mat src = imread(img_path);
	cv::Mat img = src.clone();

	// 模型初始化
	Yolov8SegOnnx task_segment_ort;

	// 模型推理
	// yolov8_onnx(task_segment_ort, img, model_Path, isGpu, output_path); //yolov8 onnxruntime segment
	yolov8_onnx_stream(task_segment_ort, img, modelBinary, isGpu, output_path);
#endif
	return 0;
}
